import pandas as pd
import numpy as np
from typing import Tuple, Dict, Any
import yaml
import os


class DataLoader:
    def __init__(self, config_path: str):
        """
        初始化数据加载器
        :param config_path: 配置文件路径
        """
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)

        self.data_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
                                      'data', 'raw')

    def load_data(self, file_name: str) -> pd.DataFrame:
        """
        加载数据
        :param file_name: 文件名
        :return: 数据框
        """
        file_path = os.path.join(self.data_path, file_name)
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"数据文件不存在: {file_path}")

        return pd.read_csv(file_path)

    def save_processed_data(self, data: pd.DataFrame, file_name: str):
        """
        保存处理后的数据
        :param data: 数据框
        :param file_name: 文件名
        """
        file_path = os.path.join(self.data_path, file_name)
        data.to_csv(file_path, index=False)

    def split_data(self, data: pd.DataFrame) -> Tuple[pd.DataFrame, pd.Series]:
        """
        分割特征和目标
        :param data: 数据框
        :return: 特征和目标
        """
        target_col = self.config['data']['target_column']

        X = data.drop(columns=[target_col])
        y = data[target_col]

        return X, y

    def validate_data(self, data: pd.DataFrame) -> Dict[str, Any]:
        """
        验证数据质量
        :param data: 数据框
        :return: 数据质量报告
        """
        report = {
            'shape': data.shape,
            'missing_values': data.isnull().sum().to_dict(),
            'duplicates': data.duplicated().sum(),
            'dtypes': data.dtypes.to_dict(),
            'numeric_stats': data.describe().to_dict() if data.select_dtypes(include=[np.number]).shape[1] > 0 else None
        }

        return report

    def preprocess_data(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        数据预处理
        :param data: 原始数据
        :return: 预处理后的数据
        """
        # 1. 处理缺失值
        data = self._handle_missing_values(data)

        # 2. 处理异常值
        data = self._handle_outliers(data)

        # 3. 处理类别特征
        data = self._process_categorical_features(data)

        return data

    def _handle_missing_values(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        处理缺失值
        :param data: 数据框
        :return: 处理后的数据框
        """
        # 数值型特征用中位数填充
        numeric_cols = self.config['data']['numeric_columns']
        for col in numeric_cols:
            if col in data.columns:
                data[col] = data[col].fillna(data[col].median())

        # 类别型特征用众数填充
        categorical_cols = self.config['data']['categorical_columns']
        for col in categorical_cols:
            if col in data.columns:
                data[col] = data[col].fillna(data[col].mode()[0])

        return data

    def _handle_outliers(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        处理异常值
        :param data: 数据框
        :return: 处理后的数据框
        """
        numeric_cols = self.config['data']['numeric_columns']

        for col in numeric_cols:
            if col in data.columns:
                # 计算四分位数
                Q1 = data[col].quantile(0.25)
                Q3 = data[col].quantile(0.75)
                IQR = Q3 - Q1

                # 定义异常值界限
                lower_bound = Q1 - 1.5 * IQR
                upper_bound = Q3 + 1.5 * IQR

                # 将异常值替换为界限值
                data[col] = data[col].clip(lower_bound, upper_bound)

        return data

    def _process_categorical_features(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        处理类别特征
        :param data: 数据框
        :return: 处理后的数据框
        """
        categorical_cols = self.config['data']['categorical_columns']

        for col in categorical_cols:
            if col in data.columns:
                # 将类别特征转换为字符串类型
                data[col] = data[col].astype(str)

        return data 